#include "RLPlanningEngine.h"
#include "RLRewardEngine.h"
#include <random>
#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <tuple>
#include <fstream>
#include <iomanip>
#include <cassert>
#include <numeric> // Required for std::accumulate
#include <cmath>   // Required for std::sqrt

namespace algo
{

    RLPlanningEngine::RLPlanningEngine(const RLPlatformData& platformData, const std::map<std::string, darr::DimsArray<float>>& topographyData, const algobase::RLConfig& config)
        : mPlatformData(platformData), mDepthData(topographyData.at("depth")), mConfig(config),
          mGridSizeX(0), mGridSizeY(0), mGridSizeZ(50),
          mCurrentEpsilon(config.epsilon_start),
          mTrainingDuration(0), mRandGen(std::random_device{}())
    {
        mRewardEngine = std::make_unique<RLRewardEngine>(mPlatformData, mConfig, *this);
#ifdef RL_DEBUG
        std::cout << "RLPlanningEngine created from map." << std::endl;
#endif
    }

    RLPlanningEngine::RLPlanningEngine(const RLPlatformData& platformData, const darr::DimsArray<float>& depthData, const algobase::RLConfig& config)
        : mPlatformData(platformData), mDepthData(depthData), mConfig(config),
        mGridSizeX(0), mGridSizeY(0), mGridSizeZ(50),
        mCurrentEpsilon(config.epsilon_start),
        mTrainingDuration(0), mRandGen(std::random_device{}())
    {
        mRewardEngine = std::make_unique<RLRewardEngine>(mPlatformData, mConfig, *this);
#ifdef RL_DEBUG
        std::cout << "RLPlanningEngine created from DimsArray." << std::endl;
#endif
    }

    RLPlanningEngine::~RLPlanningEngine()
    {
    }

    double RLPlanningEngine::findSafeDepth(const algobase::Coord& pos) const
    {
        double seabed_depth = getDepth(pos[0], pos[1]); // This will be negative
        // Start searching from a depth just below the minimum safe surface distance
        double safe_depth = -(2 * mResolutionZ + 1.0) - mResolutionZ;

        while (safe_depth > seabed_depth)
        {
            // Check if this depth is away from the seabed danger zone
            bool is_safe_from_seabed = safe_depth > seabed_depth + (2 * mResolutionZ);
            if (is_safe_from_seabed)
            {
                return safe_depth;
            }
            safe_depth -= mResolutionZ; // Check deeper (more negative)
        }

        // Fallback to a mid-point if the water is too shallow for a "safe" channel
        return (seabed_depth / 2.0);
    }

    algobase::Path RLPlanningEngine::findPath(const algobase::Coord& start, const algobase::Coord& end, const std::string& output_dir)
    {
        mOutputDir = output_dir;
        auto start_time = std::chrono::high_resolution_clock::now();

        discretizeEnvironment(start, end);
        initializeQTable();

        // --- Find Safe Starting and Ending Depths ---
        double start_depth = findSafeDepth(start);
        double end_depth = findSafeDepth(end);

        auto start_grid = worldToGrid({ start[0], start[1], start_depth });
        auto end_grid = worldToGrid({ end[0], end[1], end_depth });

#ifdef RL_DEBUG
        // --- DIAGNOSTIC LOGS ---
        std::cout << "\n--- Environment Discretization & Coordinate Conversion ---" << std::endl;
        std::cout << "World Start: (" << start[0] << ", " << start[1] << "), Safe Depth: " << start_depth << std::endl;
        std::cout << "World End:   (" << end[0] << ", " << end[1] << "), Safe Depth: " << end_depth << std::endl;
        std::cout << "Grid Size:   (" << mGridSizeX << ", " << mGridSizeY << ")" << std::endl;
        std::cout << "Resolution:  (" << mResolutionX << ", " << mResolutionY << ", " << mResolutionZ << ")" << std::endl;
        std::cout << "Grid Start:  (" << std::get<0>(start_grid) << ", " << std::get<1>(start_grid) << ", " << std::get<2>(start_grid) << ")" << std::endl;
        std::cout << "Grid End:    (" << std::get<0>(end_grid) << ", " << std::get<1>(end_grid) << ", " << std::get<2>(end_grid) << ")" << std::endl;
        
        auto recovered_start = gridToWorld(std::get<0>(start_grid), std::get<1>(start_grid), std::get<2>(start_grid));
        auto recovered_end = gridToWorld(std::get<0>(end_grid), std::get<1>(end_grid), std::get<2>(end_grid));
        std::cout << "Recovered World Start: (" << recovered_start[0] << ", " << recovered_start[1] << ", " << recovered_start[2] << ")" << std::endl;
        std::cout << "Recovered World End:   (" << recovered_end[0] << ", " << recovered_end[1] << ", " << recovered_end[2] << ")" << std::endl;
        std::cout << "--------------------------------------------------------\n" << std::endl;
#endif


#ifdef RL_LOG_TO_FILE
        std::ofstream rl_params_log(mOutputDir + "/rl_params.csv");
        rl_params_log << "episode,epsilon\n";

        std::ofstream timing_log(mOutputDir + "/timing.csv");
        timing_log << "episode,step,duration_us\n";
#endif

        for (int i = 0; i < mConfig.num_episodes; ++i)
        {
            auto current_grid = start_grid;
            Action last_action = Action::Invalid; // Initialize with an invalid action
            bool goal_reached = false;
            for (int j = 0; j < mConfig.max_steps_per_episode; ++j)
            {
#ifdef RL_DEBUG
                auto step_start_time = std::chrono::high_resolution_clock::now();
#endif

                int ix = std::get<0>(current_grid);
                int iy = std::get<1>(current_grid);
                int iz = std::get<2>(current_grid);

                Action action = chooseAction(ix, iy, iz);

                int next_ix = ix, next_iy = iy, next_iz = iz;
                switch (action)
                {
                    case North: next_iy++; break;
                    case South: next_iy--; break;
                    case East:  next_ix++; break;
                    case West:  next_ix--; break;
                    case Up:    next_iz--; break; // Depth decreases (less negative)
                    case Down:  next_iz++; break; // Depth increases (more negative)
                    case Invalid: break; // Should not happen
                }

                // Check boundaries
                if (next_ix < 0 || next_ix >= mGridSizeX || next_iy < 0 || next_iy >= mGridSizeY || next_iz < 0 || next_iz >= mGridSizeZ)
                {
                    continue; // Invalid move
                }

                // --- Final safety check before updating state ---
                assert(next_ix >= 0 && next_ix < mGridSizeX);
                assert(next_iy >= 0 && next_iy < mGridSizeY);
                assert(next_iz >= 0 && next_iz < mGridSizeZ);

                std::vector<double> reward = mRewardEngine->calculateReward(next_ix, next_iy, next_iz, static_cast<int>(last_action), static_cast<int>(action), end_grid);

                if (next_ix == std::get<0>(end_grid) && next_iy == std::get<1>(end_grid) && next_iz == std::get<2>(end_grid))
                {
                    reward[REWARD_GOAL_APPROACH] = 100.0; // Goal reward
                    goal_reached = true;
                }

                updateQValue(ix, iy, iz, action, reward, next_ix, next_iy, next_iz);
                current_grid = std::make_tuple(next_ix, next_iy, next_iz);
                last_action = action; // Update last action

#ifdef RL_LOG_TO_FILE
                auto step_end_time = std::chrono::high_resolution_clock::now();
                auto step_duration = std::chrono::duration_cast<std::chrono::microseconds>(step_end_time - step_start_time);
                timing_log << i << "," << j << "," << step_duration.count() << "\n";
#endif

                if (goal_reached)
                {
#ifdef RL_DEBUG
                    if (i % 100 == 0) {
                        std::cout << "Episode " << i << ": Goal reached in " << j << " steps." << std::endl;
                    }
#endif
                    break;
                }
            }

            // Epsilon decay
            if (mCurrentEpsilon > mConfig.epsilon_min) {
                mCurrentEpsilon *= mConfig.epsilon_decay;
            }
#ifdef RL_LOG_TO_FILE
            rl_params_log << i << "," << mCurrentEpsilon << "\n";
#endif

            // Save intermediate results every 500 episodes
            if (i > 0 && i % 500 == 0) {
#ifdef RL_DEBUG
                std::cout << "Episode " << i << ", Epsilon: " << mCurrentEpsilon << std::endl;
#endif
                saveIntermediatePath(i, start, end);
                saveQTable(i);
            }
        }
#ifdef RL_LOG_TO_FILE
        rl_params_log.close();
        timing_log.close();
#endif

        // Extract path
        algobase::Path path;
        auto current_grid = start_grid;
        algobase::PathCoord start_path_coord;
        start_path_coord.mCoordZ = gridToWorld(std::get<0>(current_grid), std::get<1>(current_grid), std::get<2>(current_grid));
        path.push_back(start_path_coord);
        
        int max_path_steps = mGridSizeX * mGridSizeY; // A safe upper bound for path length
        for (int i = 0; i < max_path_steps; ++i)
        {
            int ix = std::get<0>(current_grid);
            int iy = std::get<1>(current_grid);
            int iz = std::get<2>(current_grid);

            // Choose best action (no exploration)
            const auto& q_actions = mQTable[ix][iy][iz];
            auto best_action_it = std::max_element(q_actions.begin(), q_actions.end(), 
                [](const std::vector<double>& a, const std::vector<double>& b) {
                    return std::accumulate(a.begin(), a.end(), 0.0) < std::accumulate(b.begin(), b.end(), 0.0);
                });
            Action best_action = static_cast<Action>(std::distance(q_actions.begin(), best_action_it));
            
#ifdef RL_DEBUG
            double best_q_sum = std::accumulate(best_action_it->begin(), best_action_it->end(), 0.0);
            std::cout << "Step " << i << ": Pos(" << ix << "," << iy << "," << iz
                      << "), Best Action: " << best_action
                      << ", Q-vector Sum: " << best_q_sum << std::endl;
#endif

            int next_ix = ix, next_iy = iy, next_iz = iz;
            switch (best_action)
            {
                case North: next_iy++; break;
                case South: next_iy--; break;
                case East:  next_ix++; break;
                case West:  next_ix--; break;
                case Up:    next_iz--; break;
                case Down:  next_iz++; break;
                case Invalid: break; // Should not happen
            }

            if (next_ix < 0 || next_ix >= mGridSizeX || next_iy < 0 || next_iy >= mGridSizeY || next_iz < 0 || next_iz >= mGridSizeZ)
            {
                break; // Invalid move
            }

            current_grid = std::make_tuple(next_ix, next_iy, next_iz);
            algobase::PathCoord next_path_coord;
            next_path_coord.mCoordZ = gridToWorld(next_ix, next_iy, next_iz);
            path.push_back(next_path_coord);

            if (next_ix == std::get<0>(end_grid) && next_iy == std::get<1>(end_grid) && next_iz == std::get<2>(end_grid))
            {
                break; // Reached goal
            }
        }

        // --- Save final path ---
#ifdef RL_DEBUG
#ifdef RL_LOG_TO_FILE
        std::ofstream final_path_file(mOutputDir + "/final_path.txt");
        if (final_path_file.is_open()) {
            final_path_file << std::fixed << std::setprecision(6);
            for (const auto& p : path) {
                final_path_file << p.mCoordZ[0] << " " << p.mCoordZ[1] << " " << p.mCoordZ[2] << "\n";
            }
            final_path_file.close();
        }
#endif
#endif

        auto end_time = std::chrono::high_resolution_clock::now();
        mTrainingDuration = end_time - start_time;

        return path;
    }

    void RLPlanningEngine::discretizeEnvironment(const algobase::Coord& /*start*/, const algobase::Coord& /*end*/)
    {
        // Use the topography data to define the grid boundaries and resolution
        auto dim_names = mDepthData.dimensionNames();
        if (dim_names.size() < 2) {
            throw std::runtime_error("Topography data has fewer than 2 dimensions.");
        }
        // Assuming the order is lat, lon as per the diagnostic output
        const auto& lat_coords = mDepthData.dimensionScale(dim_names[0]);
        const auto& lon_coords = mDepthData.dimensionScale(dim_names[1]);

        if (lon_coords.size() < 2 || lat_coords.size() < 2)
        {
            throw std::runtime_error("Topography data has insufficient dimensions.");
        }

        // Set grid size from the config, not from the data resolution
        mGridSizeX = mConfig.grid_size_x;
        mGridSizeY = mConfig.grid_size_y;

        double min_lon = lon_coords.front();
        double max_lon = lon_coords.back();
        double min_lat = lat_coords.front();
        double max_lat = lat_coords.back();

        mOrigin = { min_lon, min_lat };
        
        // Calculate resolution based on the desired grid size and the geographical extent
        mResolutionX = (max_lon - min_lon) / (mGridSizeX > 1 ? mGridSizeX - 1 : 1);
        mResolutionY = (max_lat - min_lat) / (mGridSizeY > 1 ? mGridSizeY - 1 : 1);
        mResolutionZ = 500.0 / (mGridSizeZ > 1 ? mGridSizeZ - 1 : 1); // Assuming max depth of 500m
    }

    algobase::CoordZ RLPlanningEngine::gridToWorld(int ix, int iy, int iz) const
    {
        double lon = mOrigin[0] + ix * mResolutionX;
        double lat = mOrigin[1] + iy * mResolutionY;
        // Depth is negative, iz=0 is surface, iz increases -> depth increases (more negative)
        double depth = - (iz * mResolutionZ);
        return { lon, lat, depth };
    }

    std::tuple<int, int, int> RLPlanningEngine::worldToGrid(const algobase::CoordZ& pos) const
    {
        int ix = static_cast<int>(std::round((pos[0] - mOrigin[0]) / mResolutionX));
        int iy = static_cast<int>(std::round((pos[1] - mOrigin[1]) / mResolutionY));
        // World depth is negative, grid iz is positive
        int iz = static_cast<int>(std::round(-pos[2] / mResolutionZ));

        // Clamp values to be within [0, gridSize - 1]
        ix = std::max(0, std::min(ix, mGridSizeX - 1));
        iy = std::max(0, std::min(iy, mGridSizeY - 1));
        iz = std::max(0, std::min(iz, mGridSizeZ - 1));

        return std::make_tuple(ix, iy, iz);
    }

    double RLPlanningEngine::getDepth(double lon, double lat) const
    {
        auto dim_names = mDepthData.dimensionNames();
        if (dim_names.size() < 2) {
            return -500.0; // Should not happen if discretizeEnvironment succeeded
        }
        const auto& lat_coords = mDepthData.dimensionScale(dim_names[0]);
        const auto& lon_coords = mDepthData.dimensionScale(dim_names[1]);

        if (lon_coords.empty() || lat_coords.empty()) return -500.0;

        // Handle out-of-bounds cases by clamping to the edge
        if (lon <= lon_coords.front()) lon = lon_coords.front();
        if (lon >= lon_coords.back()) lon = lon_coords.back() - 1e-6; // Prevent being exactly on the last point
        if (lat <= lat_coords.front()) lat = lat_coords.front();
        if (lat >= lat_coords.back()) lat = lat_coords.back() - 1e-6;

        auto it_lon = std::upper_bound(lon_coords.begin(), lon_coords.end(), lon);
        int ix2 = static_cast<int>(std::distance(lon_coords.begin(), it_lon));
        int ix1 = ix2 - 1;

        auto it_lat = std::upper_bound(lat_coords.begin(), lat_coords.end(), lat);
        int iy2 = static_cast<int>(std::distance(lat_coords.begin(), it_lat));
        int iy1 = iy2 - 1;

        // Bilinear interpolation
        double x1 = lon_coords[ix1];
        double x2 = lon_coords[ix2];
        double y1 = lat_coords[iy1];
        double y2 = lat_coords[iy2];

        float q11 = mDepthData.at({ {dim_names[1], x1}, {dim_names[0], y1} });
        float q12 = mDepthData.at({ {dim_names[1], x1}, {dim_names[0], y2} });
        float q21 = mDepthData.at({ {dim_names[1], x2}, {dim_names[0], y1} });
        float q22 = mDepthData.at({ {dim_names[1], x2}, {dim_names[0], y2} });

        if (x2 == x1 || y2 == y1) {
            return q11; // Avoid division by zero
        }

        double r1 = ((x2 - lon) / (x2 - x1)) * q11 + ((lon - x1) / (x2 - x1)) * q21;
        double r2 = ((x2 - lon) / (x2 - x1)) * q12 + ((lon - x1) / (x2 - x1)) * q22;

        return ((y2 - lat) / (y2 - y1)) * r1 + ((lat - y1) / (y2 - y1)) * r2;
    }

    void RLPlanningEngine::initializeQTable()
    {
        mQTable.resize(mGridSizeX);
        for (int i = 0; i < mGridSizeX; ++i)
        {
            mQTable[i].resize(mGridSizeY);
            for (int j = 0; j < mGridSizeY; ++j)
            {
                mQTable[i][j].resize(mGridSizeZ);
                for (int k = 0; k < mGridSizeZ; ++k)
                {
                    for(int a = 0; a < NUM_ACTIONS; ++a)
                    {
                        mQTable[i][j][k][a].resize(REWARD_VECTOR_SIZE, 0.0);
                    }
                }
            }
        }
        // Starting reward bonus
        auto start_grid = worldToGrid({ 0, 0, findSafeDepth({0,0}) }); // A bit of a hack, but gets a safe grid
        int iz = std::get<2>(start_grid);
        if (iz >= 0 && iz < mGridSizeZ) {
            for (int i = 0; i < mGridSizeX; ++i) {
                for (int j = 0; j < mGridSizeY; ++j) {
                    for (int a = 0; a < NUM_ACTIONS; ++a) {
                        mQTable[i][j][iz][a][REWARD_SAFETY] = 10.0; // Give a high initial safety reward for the safe depth layer
                    }
                }
            }
        }
    }

    RLPlanningEngine::Action RLPlanningEngine::chooseAction(int ix, int iy, int iz)
    {
        std::uniform_real_distribution<> dis(0.0, 1.0);

        if (dis(mRandGen) < mCurrentEpsilon)
        {
            // Explore: choose a random action
            return static_cast<Action>(std::uniform_int_distribution<>(0, NUM_ACTIONS - 1)(mRandGen));
        }
        else
        {
            // Exploit: choose the best action based on the sum of Q-vector components
            const auto& q_actions = mQTable[ix][iy][iz];
            auto best_action_it = std::max_element(q_actions.begin(), q_actions.end(), 
                [](const std::vector<double>& a, const std::vector<double>& b) {
                    return std::accumulate(a.begin(), a.end(), 0.0) < std::accumulate(b.begin(), b.end(), 0.0);
                });
            return static_cast<Action>(std::distance(q_actions.begin(), best_action_it));
        }
    }

    void RLPlanningEngine::updateQValue(int ix, int iy, int iz, Action action, const std::vector<double>& reward, int next_ix, int next_iy, int next_iz)
    {
        const auto& next_q_actions = mQTable[next_ix][next_iy][next_iz];
        auto best_next_action_it = std::max_element(next_q_actions.begin(), next_q_actions.end(), 
            [](const std::vector<double>& a, const std::vector<double>& b) {
                return std::accumulate(a.begin(), a.end(), 0.0) < std::accumulate(b.begin(), b.end(), 0.0);
            });
        const std::vector<double>& max_next_q_vector = *best_next_action_it;

        std::vector<double>& old_q_vector = mQTable[ix][iy][iz][action];
        for (size_t i = 0; i < REWARD_VECTOR_SIZE; ++i)
        {
            double old_q = old_q_vector[i];
            double new_q = old_q + mConfig.alpha * (reward[i] + mConfig.gamma * max_next_q_vector[i] - old_q);
            old_q_vector[i] = new_q;
        }
    }

    double RLPlanningEngine::getTrainingDuration() const
    {
        return mTrainingDuration.count();
    }

    void RLPlanningEngine::saveIntermediatePath(int episode, const algobase::Coord& start, const algobase::Coord& end) const
    {
#ifdef RL_LOG_TO_FILE
        std::string filename = mOutputDir + "/path_ep" + std::to_string(episode) + ".txt";
        std::ofstream outfile(filename);
        if (!outfile.is_open()) {
#ifdef RL_DEBUG
            std::cerr << "Error: Could not open file for intermediate path: " << filename << std::endl;
#endif
            return;
        }
#ifdef RL_DEBUG
        std::cout << "Saving intermediate path to " << filename << std::endl;
#endif

        algobase::Path path;
        // Use the actual safe depths for saving intermediate paths
        double start_depth = findSafeDepth(start);
        double end_depth = findSafeDepth(end);
        auto start_grid = worldToGrid({ start[0], start[1], start_depth });
        auto end_grid = worldToGrid({ end[0], end[1], end_depth });
        auto current_grid = start_grid;

        int max_path_steps = mGridSizeX * mGridSizeY; // A reasonable limit
        for (int i = 0; i < max_path_steps; ++i)
        {
            int ix = std::get<0>(current_grid);
            int iy = std::get<1>(current_grid);
            int iz = std::get<2>(current_grid);

            const auto& q_actions = mQTable[ix][iy][iz];
            auto best_action_it = std::max_element(q_actions.begin(), q_actions.end(), 
                [](const std::vector<double>& a, const std::vector<double>& b) {
                    return std::accumulate(a.begin(), a.end(), 0.0) < std::accumulate(b.begin(), b.end(), 0.0);
                });
            Action best_action = static_cast<Action>(std::distance(q_actions.begin(), best_action_it));

            int next_ix = ix, next_iy = iy, next_iz = iz;
            switch (best_action)
            {
                case North: next_iy++; break;
                case South: next_iy--; break;
                case East:  next_ix++; break;
                case West:  next_ix--; break;
                case Up:    next_iz--; break;
                case Down:  next_iz++; break;
                case Invalid: break; // Should not happen
            }

            if (next_ix < 0 || next_ix >= mGridSizeX || next_iy < 0 || next_iy >= mGridSizeY || next_iz < 0 || next_iz >= mGridSizeZ)
            {
                break;
            }

            current_grid = std::make_tuple(next_ix, next_iy, next_iz);
            algobase::CoordZ p = gridToWorld(next_ix, next_iy, next_iz);
            outfile << p[0] << " " << p[1] << " " << p[2] << "\n";

            if (next_ix == std::get<0>(end_grid) && next_iy == std::get<1>(end_grid) && next_iz == std::get<2>(end_grid))
            {
                break;
            }
        }
        outfile.close();
#endif
    }

    void RLPlanningEngine::saveQTable(int episode) const
    {
#ifdef RL_LOG_TO_FILE
        std::string filename = mOutputDir + "/q_table_ep" + std::to_string(episode) + ".txt";
        std::ofstream outfile(filename);
        if (!outfile.is_open()) {
#ifdef RL_DEBUG
            std::cerr << "Error: Could not open file for Q-table: " << filename << std::endl;
#endif
            return;
        }
#ifdef RL_DEBUG
        std::cout << "Saving Q-table to " << filename << std::endl;
#endif

        outfile << std::fixed << std::setprecision(3);
        for (int k = 0; k < mGridSizeZ; ++k) {
            outfile << "Depth level: " << k << "\n";
            for (int j = 0; j < mGridSizeY; ++j) {
                for (int i = 0; i < mGridSizeX; ++i) {
                    const auto& actions = mQTable[i][j][k];
                    // Find the action with the highest summed Q-vector
                    auto best_action_it = std::max_element(actions.begin(), actions.end(), 
                        [](const std::vector<double>& a, const std::vector<double>& b) {
                            return std::accumulate(a.begin(), a.end(), 0.0) < std::accumulate(b.begin(), b.end(), 0.0);
                        });
                    double max_q_sum = std::accumulate(best_action_it->begin(), best_action_it->end(), 0.0);
                    outfile << std::setw(8) << max_q_sum;
                }
                outfile << "\n";
            }
            outfile << "\n";
        }
        outfile.close();
#endif
    }
}